﻿import os
from time import sleep
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.version
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np


class NN(nn.Module):
    def __init__(self):
        super(NN, self).__init__()
        # 定义第一层全连接层，将输入784映射到256
        self.fc1 = nn.Linear(784, 256)
        # 定义第二层全连接层，将256映射到10
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        # 首先将输入展平为一维向量
        x = x.view(x.size(0), -1)  # 假设输入是(batch_size, 1, 28, 28)
        # 通过第一层全连接层并应用ReLU激活函数
        x = torch.nn.functional.relu(self.fc1(x))
        # 通过第二层全连接层得到最终输出
        output = self.fc2(x)
        return output


transform = transforms.Compose(
    [transforms.Grayscale(num_output_channels=1), transforms.ToTensor()]
)
current_working_directory = os.getcwd()
mechine_file_path = os.path.join(
    current_working_directory,
    "Pytorch",
    "Network",
    "Nn",
    "Machine",
    "model_parameters.pth",
)

data_file_path = os.path.join(current_working_directory, "Dataset", "Minist", "File")
# 测试集
test_dataset = datasets.MNIST(
    root=data_file_path, train=False, transform=transforms.ToTensor()
)

print("test_dataset length:", len(test_dataset))
model = NN()  # 定义神经网络模型
model.load_state_dict(torch.load(mechine_file_path))  # 加载刚刚训练好的模型文件

right = 0  # 保存正确识别的数量
for i, (x, y) in enumerate(test_dataset):
    output = model(x)  # 将其中的数据x输入到模型
    predict = output.argmax(1).item()  # 选择概率最大标签的作为预测结果
    # 对比预测值predict和真实标签y
    if predict == y:
        right += 1
    else:
        image = x
        image_np = image.numpy().squeeze()
        # 显示图像
        plt.imshow(image_np, cmap='gray')
        plt.title(f'wrong case: predict ={predict} y = {y} ')
        plt.show()
        sleep(10)
       
# 计算出测试效果
sample_num = len(test_dataset)
acc = right * 1.0 / sample_num
print("test accuracy = %d / %d = %.3lf" % (right, sample_num, acc))